use axum::{
    async_trait,
    body::Body,
    extract::{FromRequest, Path},
    response::{IntoResponse, Response},
    RequestExt, RequestPartsExt,
};
use axum_extra::{
    headers::{authorization::Bearer, Authorization},
    TypedHeader,
};
use bytes::{BufMut, BytesMut};
use http::{Request, StatusCode};
use log::info;
use ruma::{
    api::{client::error::ErrorKind, AuthScheme, IncomingRequest, OutgoingResponse},
    CanonicalJsonValue, OwnedDeviceId, OwnedUserId, UserId,
};
use serde::Deserialize;
use tracing::{debug, error, warn};

use super::{Ruma, RumaResponse};
use crate::{
    server::{self, server},
    utils::{Error, Result},
};

enum Token {
    User((OwnedUserId, OwnedDeviceId)),
    Invalid,
    None,
}

#[async_trait]
impl<T, S> FromRequest<S> for Ruma<T>
where
    T: IncomingRequest,
{
    type Rejection = Error;

    async fn from_request(req: Request<Body>, _state: &S) -> Result<Self, Self::Rejection> {
        #[derive(Deserialize)]
        struct QueryParams {
            access_token: Option<String>,
            _user_id: Option<String>,
        }

        let (mut parts, mut body) = {
            let limited_req = req.with_limited_body();
            let (parts, body) = limited_req.into_parts();
            let body = axum::body::to_bytes(
                body,
                server::server()
                    .config
                    .max_request_size
                    .try_into()
                    .unwrap_or(usize::MAX),
            )
            .await
            .map_err(|_| Error::BadRequest(ErrorKind::MissingToken, "Missing token."))?;
            (parts, body)
        };

        let metadata = T::METADATA;
        let auth_header: Option<TypedHeader<Authorization<Bearer>>> = parts.extract().await?;
        let path_params: Path<Vec<String>> = parts.extract().await?;

        let query = parts.uri.query().unwrap_or_default();
        let query_params: QueryParams = match serde_html_form::from_str(query) {
            Ok(params) => params,
            Err(e) => {
                error!(%query, "Failed to deserialize query parameters: {}", e);
                return Err(Error::BadRequest(
                    ErrorKind::Unknown,
                    "Failed to read query parameters",
                ));
            }
        };

        let token = match &auth_header {
            Some(TypedHeader(Authorization(bearer))) => Some(bearer.token()),
            None => query_params.access_token.as_deref(),
        };

        let token = if let Some(token) = token {
            if let Some((user_id, device_id)) = server().users.find_from_token(token)? {
                Token::User((user_id, OwnedDeviceId::from(device_id)))
            } else {
                Token::Invalid
            }
        } else {
            Token::None
        };

        let mut json_body = serde_json::from_slice::<CanonicalJsonValue>(&body).ok();

        let (sender_user, sender_device, sender_servername) = match (metadata.authentication, token)
        {
            (_, Token::Invalid) => {
                // OpenID endpoint uses a query param with the same name, drop this once query params for user auth are removed from the spec
                if query_params.access_token.is_some() {
                    (None, None, None)
                } else {
                    return Err(Error::BadRequest(
                        ErrorKind::UnknownToken { soft_logout: false },
                        "Unknown access token.",
                    ));
                }
            }
            (AuthScheme::AccessToken, Token::None) => {
                return Err(Error::BadRequest(
                    ErrorKind::MissingToken,
                    "Missing access token.",
                ));
            }
            (
                AuthScheme::None | AuthScheme::AppserviceToken | AuthScheme::AccessTokenOptional,
                Token::None,
            ) => (None, None, None),
            (AuthScheme::AppserviceToken, Token::User(_)) => {
                return Err(Error::BadRequest(
                    ErrorKind::Unauthorized,
                    "Only appservice access tokens should be used on this endpoint.",
                ));
            }
            (
                AuthScheme::AccessToken | AuthScheme::AccessTokenOptional | AuthScheme::None,
                Token::User((user_id, device_id)),
            ) => (
                Some(user_id.clone()),
                Some(device_id),
                Some(user_id.server_name().to_owned()),
            ),

            (_, _) => {
                return Err(Error::BadRequest(
                    ErrorKind::Unauthorized,
                    "Access token required.",
                ));
            }
        };

        let mut http_request = Request::builder().uri(parts.uri).method(parts.method);
        *http_request.headers_mut().unwrap() = parts.headers;

        if let Some(CanonicalJsonValue::Object(json_body)) = &mut json_body {
            let _user_id = sender_user.clone().unwrap_or_else(|| {
                UserId::parse_with_server_name("", &server().config.server_name)
                    .expect("we know this is valid")
            });

            // let uiaa_request = json_body
            //     .get("auth")
            //     .and_then(|auth| auth.as_object())
            //     .and_then(|auth| auth.get("session"))
            //     .and_then(|session| session.as_str())
            //     .and_then(|session| {
            //         services().uiaa.get_uiaa_request(
            //             &user_id,
            //             &sender_device.clone().unwrap_or_else(|| "".into()),
            //             session,
            //         )
            //     });

            // if let Some(CanonicalJsonValue::Object(initial_request)) = uiaa_request {
            //     for (key, value) in initial_request {
            //         json_body.entry(key).or_insert(value);
            //     }
            // }

            let mut buf = BytesMut::new().writer();
            serde_json::to_writer(&mut buf, json_body).expect("value serialization can't fail");
            body = buf.into_inner().freeze();
        }

        let http_request = http_request.body(&*body).unwrap();

        debug!("{:?}", http_request);

        let body = T::try_from_http_request(http_request, &path_params).map_err(|e| {
            warn!("try_from_http_request failed: {:?}", e);
            debug!("JSON body: {:?}", json_body);
            Error::BadRequest(ErrorKind::BadJson, "Failed to deserialize request.")
        })?;

        Ok(Ruma {
            body,
            sender_user,
            sender_device,
            sender_servername,
            json_body,
        })
    }
}

impl<T: OutgoingResponse> IntoResponse for RumaResponse<T> {
    fn into_response(self) -> Response {
        match self.0.try_into_http_response::<BytesMut>() {
            Ok(res) => res.map(BytesMut::freeze).map(Body::from).into_response(),
            Err(_) => StatusCode::INTERNAL_SERVER_ERROR.into_response(),
        }
    }
}
